import cv2
import os


class VideoFileReader(object):
    def __init__(self, video_bytes):
        # save video locally
        self.save_path = '../detect/temp/'
        if not os.path.exists(self.save_path):
            os.makedirs(self.save_path)
        with open(self.save_path + 'video.avi', 'bw') as self.f:
            self.f.write(video_bytes)
        # init cap
        self.cap = cv2.VideoCapture(self.save_path + 'video.avi')
        self.count = 0

    def __iter__(self):
        return self

    def __next__(self):
        if not self.cap.grab():
            self.cap.release()                       # release cap
            os.remove(self.save_path + 'video.avi')  # clear local video
            raise StopIteration
        img = None
        if not self.cap.isOpened():
            self.cap.open(self.save_path + 'video.avi')
        ret, img = self.cap.read()
        if ret:
            self.count += 1
            return img
        return

    def __len__(self):
        return self.count
